package rahmath.sharif.docrypt.util;

import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.File;
import java.io.FileInputStream;
import java.io.FileOutputStream;
import java.io.IOException;
import java.io.ObjectInputStream;
import java.io.ObjectOutputStream;
import java.io.Serializable;
import java.security.SecureRandom;

import javax.crypto.Cipher;
import javax.crypto.KeyGenerator;
import javax.crypto.SecretKey;
import javax.crypto.spec.SecretKeySpec;

import android.util.Base64;

public class AESCryptor {

	private static final String UTF8 = "UTF-8";
	private static final String ALGORITHM = "AES";
	private static final String RANDOM_ALGORITHM = "SHA1PRNG";
	private byte[] rawKey;

	public AESCryptor(int bitLength, String seed) {
		try {
			rawKey = getRawKey(bitLength, seed.getBytes(UTF8));
		} catch (Exception e) {
			e.printStackTrace();
		}
	}

	public String encrypt(String cleartext) {
		byte[] result;
		try {
			result = encrypt(this.rawKey, cleartext.getBytes(UTF8));
			return Base64.encodeToString(result, Base64.DEFAULT);
		} catch (Exception e) {
			e.printStackTrace();
		}
		return null;
	}

	public String decrypt(String encrypted) {
		try {
			byte[] enc = Base64.decode(encrypted, Base64.DEFAULT);
			byte[] result = decrypt(this.rawKey, enc);
			return new String(result);
		} catch (Exception e) {
			e.printStackTrace();
		}
		return null;
	}

	public static Cipher getCipher(String seed, int mode) {
		try {
			SecretKeySpec skeySpec = new SecretKeySpec(getRawKey(256,
					seed.getBytes(UTF8)), ALGORITHM);
			Cipher cipher = Cipher.getInstance(ALGORITHM);
			switch (mode) {
			case Cipher.ENCRYPT_MODE:
				cipher.init(Cipher.ENCRYPT_MODE, skeySpec);
				return cipher;
			case Cipher.DECRYPT_MODE:
				cipher.init(Cipher.DECRYPT_MODE, skeySpec);
				return cipher;
			}
		} catch (Exception e) {
			e.printStackTrace();
		}
		return null;
	}

	public static String encrypt(int bitLength, String seed, String cleartext) {
		try {
			byte[] rawKey = getRawKey(bitLength, seed.getBytes(UTF8));
			byte[] result = encrypt(rawKey, cleartext.getBytes(UTF8));
			return Base64.encodeToString(result, Base64.DEFAULT);
		} catch (Exception e) {
			e.printStackTrace();
		}
		return null;
	}

	public static String decrypt(int bitLength, String seed, String encrypted) {
		try {
			byte[] rawKey = getRawKey(bitLength, seed.getBytes(UTF8));
			byte[] enc = Base64.decode(encrypted, Base64.DEFAULT);
			byte[] result = decrypt(rawKey, enc);
			return new String(result);
		} catch (Exception e) {
			e.printStackTrace();
		}
		return null;
	}

	public static void encrypt(String seed, File clearFile, File cipherFile) {
		try {
			FileInputStream fileInputStream = new FileInputStream(clearFile);
			FileOutputStream fileOutputStream = new FileOutputStream(cipherFile);
			SecretKeySpec skeySpec = new SecretKeySpec(getRawKey(256,
					seed.getBytes(UTF8)), ALGORITHM);
			Cipher cipher = Cipher.getInstance(ALGORITHM);
			cipher.init(Cipher.ENCRYPT_MODE, skeySpec);
			byte[] input = new byte[40960];
			int bytesRead;
			while ((bytesRead = fileInputStream.read(input)) != -1) {
				byte[] output = cipher.update(input, 0, bytesRead);
				if (output != null)
					fileOutputStream.write(output);
			}
			byte[] output = cipher.doFinal();
			if (output != null)
				fileOutputStream.write(output);
			fileInputStream.close();
			fileOutputStream.flush();
			fileOutputStream.close();
		} catch (Exception e) {
			e.printStackTrace();
		}
	}

	public static void decrypt(String seed, File cipherFile, File clearFile) {
		try {
			FileInputStream fileInputStream = new FileInputStream(cipherFile);
			FileOutputStream fileOutputStream = new FileOutputStream(clearFile);
			SecretKeySpec skeySpec = new SecretKeySpec(getRawKey(256,
					seed.getBytes(UTF8)), ALGORITHM);
			Cipher cipher = Cipher.getInstance(ALGORITHM);
			cipher.init(Cipher.DECRYPT_MODE, skeySpec);
			byte[] input = new byte[40960];
			int bytesRead;
			while ((bytesRead = fileInputStream.read(input)) != -1) {
				byte[] output = cipher.update(input, 0, bytesRead);
				if (output != null)
					fileOutputStream.write(output);
			}
			byte[] output = cipher.doFinal();
			if (output != null)
				fileOutputStream.write(output);
			fileInputStream.close();
			fileOutputStream.flush();
			fileOutputStream.close();
		} catch (Exception e) {
			e.printStackTrace();
		}
	}

	public static String encryptObject(String seed,
			Serializable serializableObject) {
		try {
			byte[] rawKey = getRawKey(256, seed.getBytes(UTF8));
			byte[] result = encrypt(rawKey, serialize(serializableObject));
			return Base64.encodeToString(result, Base64.DEFAULT);
		} catch (Exception e) {
			e.printStackTrace();
		}
		return null;
	}

	public static Serializable decryptObject(String seed, String encodedString) {
		try {
			byte[] rawKey = getRawKey(256, seed.getBytes(UTF8));
			byte[] enc = Base64.decode(encodedString, Base64.DEFAULT);
			byte[] result = decrypt(rawKey, enc);
			return deserialize(result);
		} catch (Exception e) {
			e.printStackTrace();
		}
		return null;
	}

	public static byte[] serialize(Serializable obj) {
		try {
			ByteArrayOutputStream out = new ByteArrayOutputStream();
			ObjectOutputStream os = new ObjectOutputStream(out);
			os.writeObject(obj);
			return out.toByteArray();
		} catch (IOException e) {
			e.printStackTrace();
		}
		return null;
	}

	public static Serializable deserialize(byte[] data) {
		try {
			ByteArrayInputStream in = new ByteArrayInputStream(data);
			ObjectInputStream is = new ObjectInputStream(in);
			return (Serializable) is.readObject();
		} catch (Exception e) {
			e.printStackTrace();
		}
		return null;
	}

	public static byte[] getRawKey(int bitLength, byte[] seed) {
		try {
			KeyGenerator kgen = KeyGenerator.getInstance(ALGORITHM);
			SecureRandom sr = SecureRandom.getInstance(RANDOM_ALGORITHM);
			sr.setSeed(seed);
			kgen.init(bitLength, sr);
			SecretKey skey = kgen.generateKey();
			byte[] raw = skey.getEncoded();
			return raw;
		} catch (Exception e) {
			e.printStackTrace();
		}
		return null;
	}

	private static byte[] encrypt(byte[] raw, byte[] clear) throws Exception {
		SecretKeySpec skeySpec = new SecretKeySpec(raw, ALGORITHM);
		Cipher cipher = Cipher.getInstance(ALGORITHM);
		cipher.init(Cipher.ENCRYPT_MODE, skeySpec);
		byte[] encrypted = cipher.doFinal(clear);
		return encrypted;
	}

	private static byte[] decrypt(byte[] raw, byte[] encrypted)
			throws Exception {
		SecretKeySpec skeySpec = new SecretKeySpec(raw, ALGORITHM);
		Cipher cipher = Cipher.getInstance(ALGORITHM);
		cipher.init(Cipher.DECRYPT_MODE, skeySpec);
		byte[] decrypted = cipher.doFinal(encrypted);
		return decrypted;
	}

	public static String encrypt(int bitLength, char[] newPasswordFromUser,
			String string) {
		return encrypt(bitLength, new String(newPasswordFromUser), string);
	}

	public static String decrypt(int bitLength, char[] password, String string) {
		return decrypt(bitLength, new String(password), string);
	}

	public static String encryptPassword(int bitLength, String password) {
		return encrypt(bitLength, password, password);
	}

	public static boolean validatePassword(int bitLength,
			String encryptedPassword, String clearPassword) {
		return encryptedPassword.equals(encryptPassword(bitLength,
				clearPassword));
	}

	public static void erasePassword(char[] dedicatedPassword) {
		for (@SuppressWarnings("unused")
		char c : dedicatedPassword) {
			c = '0';
		}
	}

}